# predict from any given image
# author-by: xjtu-blacksmith
# create-on: 2020.2.23

import os
from os import path

from PIL import Image
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms

from config import CLASS_NUM
from pretrain import build_data
from model import vgg16_101
from utils.translate import translate_id

state_dict_dir = path.join('output', 'model-vgg16-best_acc.pth')

def predict_image(img_path):

    input = Image.open(img_path)
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.5] * 3, [0.5] * 3)
    ])
    input = transform(input)

    net = vgg16_101()
    net.load_state_dict(torch.load(state_dict_dir))
    net.eval()

    with torch.no_grad():

        output = net(input[None, ...])
        output = F.softmax(output.data, dim=1)
        topk, topclass = output.topk(5, dim=1)
        top5classes = [translate_id(class_name) for class_name in topclass[0]]
        print('Top-5 result:', ', '.join(name for name in top5classes))
        print('Probability:', ', '.join(("%.2f%%" % (p * 100)) for p in topk[0]))
        del output

if __name__ == "__main__":

    img_folder = path.join('data', 'test')
    img_list = os.listdir(img_folder)
    for img in img_list:
        print('Predicting `%s` ...' % img)
        predict_image(path.join(img_folder, img))
